package com.zjj.lbw.ai;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.client.advisor.VectorStoreChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.evaluation.EvaluationRequest;
import org.springframework.ai.evaluation.EvaluationResponse;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.model.Media;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.generation.augmentation.QueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.expansion.QueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.join.DocumentJoiner;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor.FILTER_EXPRESSION;

@RestController
public class ChatController {

    private final ChatClient chatClient;

    public ChatController(ChatClient.Builder chatClientBuilder) {
        this.chatClient = chatClientBuilder
//                .defaultSystem("请用我提供的主题，写一首七言绝句")
                .defaultAdvisors(new MessageChatMemoryAdvisor(new InMemoryChatMemory()))
                .build();
    }

    @Resource
    private ChatModel chatModel;

    @Resource
    private EmbeddingModel embeddingModel;

    @Autowired
    private DocumentService documentService;

    @Autowired
    private VectorStore vectorStore;


    @Autowired
    private ToolCallbackProvider toolCallbackProvider;

    @GetMapping("/chat")
    public String chat() {
        return this.chatClient.prompt()
                .user("你是谁")
                .call()
                .content();
    }

    @GetMapping("/chat1")
    public String chat1() {
        // 不使用公共的ChatClient，自己之间利用ChatModel构建一个
        ChatClient client = ChatClient.builder(chatModel).build();
        return client
                .prompt()
                .user("你是谁")
                .call()
                .content();
    }

    @GetMapping("/entity")
    public ActorFilms entity() {
        return this.chatClient
                .prompt()
                .user("给我周星驰的三部电影信息")
                .call()
                .entity(ActorFilms.class);
    }

    @GetMapping("/output")
    public List<ActorFilms> output() {

        BeanOutputConverter<List<ActorFilms>> outputConverter = new BeanOutputConverter<>(
                new ParameterizedTypeReference<List<ActorFilms>>() {
                });

        String format = outputConverter.getFormat();
        String template = """
                给我周星驰和成龙的三部电影信息
                {format}
                """;

        Prompt prompt = new PromptTemplate(template, Map.of("format", format)).create();

        String content = chatClient.prompt(prompt).call().content();

        return outputConverter.convert(content);
    }

    @GetMapping("/output1")
    public List<ActorFilms> output1() {
        BeanOutputConverter<List<ActorFilms>> outputConverter = new BeanOutputConverter<>(
                new ParameterizedTypeReference<List<ActorFilms>>() {
                });

        return chatClient.prompt("给我周星驰和成龙的三部电影信息").call().entity(outputConverter);
    }

    @GetMapping(value = "/streamChat", produces = "text/html;charset=UTF-8")
    public Flux<String> streamChat() {
        return chatClient
                .prompt()
                .user("你是谁")
                .stream()
                .content();
    }

    @GetMapping(value = "/sseChat")
    public SseEmitter sseChat() {
        SseEmitter sseEmitter = new SseEmitter() {
            @Override
            protected void extendResponse(ServerHttpResponse outputMessage) {
                HttpHeaders headers = outputMessage.getHeaders();
                headers.setContentType(new MediaType("text", "event-stream", StandardCharsets.UTF_8));
            }
        };

        Flux<String> stream = chatClient
                .prompt()
                .user("你是谁")
                .stream()
                .content();

        stream.subscribe(token -> {
            try {
                sseEmitter.send(token);
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
        }, sseEmitter::completeWithError, sseEmitter::complete);
        return sseEmitter;
    }

    @GetMapping("/systemChat")
    public String systemChat() {
        return this.chatClient
                .prompt()
//                .system("请用我提供的主题，写一首七言绝句")
                .advisors(new ZjjCallAroundAdvisor())
                .user("程序员")
                .call()
                .content();
    }

    @GetMapping("/chatMemoryChat")
    public String chatMemoryChat(@RequestParam("chatId") String chatId, @RequestParam("message") String message) {
        return chatClient
                .prompt()
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .user(message)
                .call()
                .content();
    }

    @GetMapping("/multimodalChat")
    public String multimodalChat() {
        ClassPathResource imageData = new ClassPathResource("/test.jpg");

        var userMessage = new UserMessage("图片里有什么?",
                List.of(new Media(MimeTypeUtils.IMAGE_PNG, imageData)));

        return chatClient.prompt().messages(userMessage).call().content();
    }

    @GetMapping("/embedding")
    public float[] embedding(@RequestParam String message) {
        return embeddingModel.embed(message);
    }

    @GetMapping("/documentEmbedding")
    public List<Document> documentEmbedding() {
        return documentService.loadText();
    }

    @GetMapping("/documentSearch")
    public List<Document> documentSearch(@RequestParam String message) {
        return documentService.search(message);
    }

    @GetMapping("/ragChat")
    public String ragChat(@RequestParam String message) {

        // 向量搜索
        List<Document> documentList = documentService.search(message);

        // 提示词模板
        PromptTemplate promptTemplate = new PromptTemplate("{userMessage}\n\n 用以下信息回答问题，主要要去掉多余的符号:\n {contents}");

        // 组装提示词
        Prompt prompt = promptTemplate.create(Map.of("userMessage", message, "contents", documentList));

        // 调用大模型
        return chatClient.prompt(prompt).call().content();
    }

    @GetMapping("/qa1")
    public String qa1(@RequestParam String message) {
        return chatClient
                .prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .user(message)
                .call()
                .content();
    }

    @GetMapping("/qa")
    public String qa(@RequestParam("chatId") String chatId, @RequestParam String message) {
        return chatClient
                .prompt()
                .advisors(VectorStoreChatMemoryAdvisor.builder(vectorStore).build())
                .advisors(c -> c.param(FILTER_EXPRESSION, "author in ['zjj', 'jill'] && 'article_type' == 'blog'"))
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .user(message)
                .call()
                .content();
    }

    private ChatMemory chatMemory = new InMemoryChatMemory();

    @GetMapping("/rag")
    public String rag(@RequestParam("chatId") String chatId, @RequestParam String message) {

        ChatClient.Builder builder = ChatClient.builder(chatModel);

        // 给定以下对话历史记录和后续查询，您的任务是合成一个简洁、独立的查询，其中包含历史记录中的上下文。确保独立查询清晰、具体，并维护用户的意图。
        QueryTransformer queryTransformer = new CompressionQueryTransformer(builder, null);

        QueryExpander queryExpander = new MultiQueryExpander(builder, null, null, null);

        DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.4)
                .topK(5)
                .build();

        DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();

        QueryAugmenter queryAugmenter = new ContextualQueryAugmenter(null, null ,null);

        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .queryTransformers(queryTransformer)
                .queryExpander(queryExpander)
                .documentRetriever(documentRetriever)
                .documentJoiner(documentJoiner)
                .queryAugmenter(queryAugmenter)
                .build();

        return chatClient
                .prompt()
                .advisors(new MessageChatMemoryAdvisor(chatMemory), retrievalAugmentationAdvisor)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .user(message)
                .call()
                .content();
    }

    @GetMapping("/evaluation")
    public EvaluationResponse evaluation(@RequestParam String message) {

        // 向量搜索
        List<Document> documentList = documentService.search(message);

        // 提示词模板
        PromptTemplate promptTemplate = new PromptTemplate("{userMessage}\n\n 用以下信息回答问题:\n {contents}");

        // 组装提示词
        Prompt prompt = promptTemplate.create(Map.of("userMessage", message, "contents", documentList));

        // 调用大模型
        String result = chatClient.prompt(prompt).call().content();
//        String result = "我是周瑜";

        var relevancyEvaluator = new RelevancyEvaluator(ChatClient.builder(chatModel));

        // 评估是否产生了幻觉
        EvaluationRequest evaluationRequest = new EvaluationRequest(message, documentList, result);

        EvaluationResponse evaluationResponse = relevancyEvaluator.evaluate(evaluationRequest);

        return evaluationResponse;
    }

    @GetMapping("/toolChat")
    public String toolChat(@RequestParam("message") String message) {
        return chatClient
                .prompt()
                .user(message)
                .tools(new ZjjTools())
                .call()
                .content();
    }

    @GetMapping("/userControlledToolChat")
    public String userControlledToolChat(@RequestParam("message") String message) {

        ToolCallingManager toolCallingManager = ToolCallingManager.builder().build();

        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(ToolCallbacks.from(new ZjjTools()))
                .internalToolExecutionEnabled(false) // 设置为false 内部不会递归去执行工具
                .build();
        Prompt prompt = new Prompt(message, chatOptions);

        ChatResponse chatResponse = chatModel.call(prompt);

        while (chatResponse.hasToolCalls()) {
            ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);

            prompt = new Prompt(toolExecutionResult.conversationHistory(), chatOptions);

            chatResponse = chatModel.call(prompt);
        }

        return chatResponse.getResult().getOutput().getText();
    }

    @GetMapping("/mcpChat")
    public String mcpChat(@RequestParam("message") String message) {
        return chatClient
                .prompt()
                .user(message)
                .tools(toolCallbackProvider.getToolCallbacks())
                .call()
                .content();
    }

}
